#==============================================================================
#
"""
    Simple memory database implementation.
"""
#
#   Copyright (C) 2007 Michael A. Muller
#   Portions Copyright (C) 2008 Google, Inc.
#
#   This file is part of ODB.
#
#   ODB is free software: you can redistribute it and/or modify it under the 
#   terms of the GNU Lesser General Public License as published by the Free 
#   Software Foundation, either version 3 of the License, or (at your option) 
#   any later version.
#
#   ODB is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU Lesser General Public License for more details.
#
#   You should have received a copy of the GNU Lesser General Public License 
#   along with ODB.  If not, see <http://www.gnu.org/licenses/>.
#
#==============================================================================

import threading, os, sys, traceback, fcntl
from pickle import Pickler, Unpickler
from odbi import *
from weakref import WeakKeyDictionary, ref

import struct
from cStringIO import StringIO

LOG_PATTERN = 'log.%09d'

class _Action:
    """
        An action (Command pattern implementation)
    """

    def doIt(self, env):
        "Apply the action to the database environment"
        raise NotImplementedError()

    def undoIt(self, env):
        "Undo the action on the environment"
        raise NotImplementedError()

    def dump(self, out, indent = ''):
        out.write('%s%s {\n' % 
                   (indent, self.__class__.__name__))
        for key, val in self.__dict__.items():
            out.write('%s  %s = %s\n' % (indent, key, repr(val)))
        out.write('%s}\n' % indent)

class _CreateBTreeDBAction(_Action):
    "Action to create a BTree table"

    def __init__(self, name, allowDups):
        self.name = name
        self.allowDups = allowDups

    def doIt(self, env):
        env._tables[self.name] = env._createBTreeDB(self.name, self.allowDups)

    def undoIt(self, env):
        try:
            del env._tables[self.name]
        except KeyError:
            pass

class _CreateBListDBAction(_Action):
    "Action to create a BList table"
    
    def __init__(self, name):
        self.name = name

    def doIt(self, env):
        env._tables[self.name] = env._createBListDB(self.name)

    def undoIt(self, env):
        try:
            del env._tables[self.name]
        except KeyError:
            pass

class _InsertAction(_Action):

    def __init__(self, name, key, val, oldVal, gotOldVal):
        self.name = name
        self.key = key
        self.val = val
        self.oldVal = oldVal
        self.gotOldVal = gotOldVal

    def __setstate__(self, d):
        self.__dict__ = d

    def doIt(self, env):
        env._tables[self.name]._insert(self.key, self.val)
    
    def undoIt(self, env):
        db = env._tables[self.name]
        if self.gotOldVal:
            db._insert(self.key, self.oldVal)
        else:
            db._delete(self.key)

class _ReplaceAction(_Action):

    def __init__(self, name, key, val, oldVal, gotOldVal):
        self.name = name
        self.key = key
        self.val = val
        self.oldVal = oldVal
        self.gotOldVal = gotOldVal

    def __setstate__(self, d):
        self.__dict__ = d

    def doIt(self, env):
        env._tables[self.name]._replace(self.key, self.val)
    
    def undoIt(self, env):
        db = env._tables[self.name]
        if self.gotOldVal:
            db._replace(self.key, self.oldVal)
        else:
            db._delete(self.key)

class _DeleteAction(_Action):
    def __init__(self, name, key, oldVal):
        self.name = name
        self.key = key
        self.oldVal = oldVal

    def doIt(self, env):
        env._tables[self.name]._delete(self.key)

    def undoIt(self, env):
        env._tables[self.name]._insert(self.key, self.oldVal)

class _PopAction(_Action):
    def __init__(self, name, index, oldVal):
        self.name = name
        self.index = index
        self.oldVal = oldVal
    
    def doIt(self, env):
        env._tables[self.name]._pop(self.index)
    
    def undoIt(self, env):
        env._tables[self.name]._insert(self.index, self.oldVal)

class _AppendAction(_Action):

    def __init__(self, name, val):
        self.name = name
        self.val = val
    
    def doIt(self, env):
        env._tables[self.name]._append(self.val)
    
    def undoIt(self, env):
        table = env._tables[self.name]
        assert table._pop() == self.val

class _CreateHashDBAction(_Action):
    "Action to create a hash db"

    def __init__(self, name, allowDups):
        self.name = name
        self.allowDups = allowDups

    def doIt(self, env):
        env._tables[self.name] = MemHashDB(env, self.name)

    def undoIt(self, env):
        try:
            del env._tables[self.name]
        except KeyError:
            pass

def _needTxn(func, env, txn, *args):
    """
        Wraps functions that need to run in a transaction.
        
        Usually these get passed a transaction as an optional argument.  This 
        checks the transaction, if it is None it creates a temporary 
        transaction and commits or aborts it depending on the outcome.
        
        parms:
            func: [callable<txn, ...>] the function that contains the 
                transaction processing.
            env: [Env]
            txn: [Txn or None]
            *args: list of additional args to be passed to func().
    """
    # create a temporary transaction if none exists
    if not txn:
        myTxn = txn = env.getTxn()
    else:
        assert txn.isActiveOn(env)
        myTxn = None
    
    try:
        result = func(txn, *args)
        if myTxn:
            myTxn.commit()
            myTxn = None
        return result
    finally:
        if myTxn: myTxn.abort()


class DoubleLockError(Exception):
    """
        Raised when an attempt is made to lock an environment that is already
        locked.  This should only happen when you create more than one
        transaction on the same environment.
    """
    pass

class SafeLock:
    """
        Lock implementation that raises an exception if it is locked more than
        once from the same thread.
    """

    def __init__(self):
        self.__impl = threading.RLock()
        self.__counter = 0
        self.__traceback = None

    def acquire(self):
        self.__impl.acquire()
        if self.__counter:
            self.__impl.release()
            raise DoubleLockError('Double locked database environment.  '
                                   'Previous lock acquired here:' +
                                   ''.join(self.__traceback)
                                  )
        self.__traceback = traceback.format_stack()
        self.__counter += 1

    def release(self):
        assert self.__counter > 0
        self.__counter -= 1
        self.__traceback = ['NOT LOCKED!']
        self.__impl.release()

class Filer:
    """
        Abstract interface to most of the file system.  This allows 
        application code to do things like automatic backups and changing 
        where logfiles and statefiles are stored.
    """
    
    # file open modes.
    READONLY = 0
    WRITABLE = 1
    CREATE = 2
    
    def init(self):
        """
            Initialize the filer.  This should create directories necessary to 
            support the creation of log and state files.
            
            Must be implemented by derived classes.
        """
        raise NotImplementedError()

    def openLogfile(self, name, mode):
        """
            Open a logfile for read or write.  Returns a standard file object.
            
            Must be implemented by derived classes.
            
            parms:
                name: [str] the unqulified logfile name.
                mode: [int] There are three values for this:
                    READONLY - open the file read-only 
                    WRITABLE - open the file for read and write
                    CREATE - create/truncate the file, opening for read and 
                        write.
        """
        raise NotImplementedError()

    def getAllLogfiles(self):
        """
            Returns the list of all logfiles (list<str>) sorted by name.
            
            Must be implemented by derived classes.
        """
        raise NotImplementedError()

    def getAllStateFiles(self):
        """
            Returns the list all state file names (list<str>) sorted by name.
            
            Must be implemented by derived classes
        """
        raise NotImplementedError()
    
    def openStateFile(self, name, mode):
        """
            Returns a file object for the state file.
            
            Must be implemented by derived classes.
            
            parms:
                name: [str] the unqualified state file name.
                mode: [int] See @openLogfile().parms.mode
        """
        raise NotImplementedError()
    
    def getStateFilePath(self, name):
        """
            If 'name' is stored on an accessible filesystem, returns the full 
            path of the state file.  Otherwise returns None.
            
            This method is provided so that backup utilities can make faster 
            copies of state files than is possible by a block-by-block copy 
            done using a statefile opened with @openStateFile().

            Must be implemented by derived classes.

            parms:
                name: [str] state file name.
        """
        raise NotImplementedError()

class DefaultFiler(Filer):
    
    def __init__(self, logPath, statePath):
        self.__logPath = logPath
        self.__statePath = statePath
        self.__stateFiles = set()

    def init(self):
        "Implements @Filer.init()"
        for path in (self.__logPath, self.__statePath):
            if not os.path.exists(path):
                os.mkdir(path)
                # store a marker indicating that this is a mem database
                file(os.path.join(path, '.mem'), 'w')

    def _openFile(self, base, name, mode):
        """
            Opens a file with the given mode.
            
            parms: (see @openLogfile().parms)
        """
        try:
            mode = {
                Filer.READONLY: 'rb',
                Filer.WRITABLE: 'rb+',
                Filer.CREATE: 'wb+'
            }[mode]
        except KeyError:
            assert False, 'Invalid file mode: %s' % mode

        return file(os.path.join(base, name), mode)

    def openLogfile(self, name, mode):
        "Implements @Filer.openLogfile()"
        return self._openFile(self.__logPath, name, mode)

    def getAllLogfiles(self):
        "Implements @Filer.getAllLogfiles()"
        files = os.listdir(self.__logPath)
        files.sort()
        return [f for f in files if f.startswith('log.')]

    def getAllStateFiles(self):
        "Implements @Filer.getAllStateFiles()"
        files = []
        for file in os.listdir(self.__statePath):
            if not file.startswith('log.'):
                files.append(file)
        return files
#        return sorted(self.__stateFiles)

    def openStateFile(self, name, mode):
        "Implements @Filer.openLogfile()"
        fp = self._openFile(self.__statePath, name, mode)
        self.__stateFiles.add(name)
        return fp

    def getStateFilePath(self, name):
        "Implements @Filer.getStateFilePath()"
        return os.path.join(self.__statePath, name)

class MemEnv(Env):

    DEFAULT_MAX_TXN_FILE_SIZE = 100000
    def __loadTxn(self):
        data = self.__log.read(4)
        if not data:
            return None
        size = struct.unpack('I', data)[0]
        data = self.__log.read(size)
        return Unpickler(StringIO(data)).load()

    def __dumpTxn(self, txn):
        dst = StringIO()
        pickler = Pickler(dst)
        pickler.dump(txn)
        data = dst.getvalue()
        self.__log.write(struct.pack('I', len(data)))
        self.__log.write(data)
        pos = self.__log.tell()
        if pos > self.__maxTxnFileSize:
            self.__checkpointAndCallback()
            self.__openNewLog()

    def __readLog(self):
        """
            Read all outstanding transactions in the logfile.
        """
        while True:
            txn = self.__loadTxn()
            if not txn:
                break
            txn._apply(self)
#                    try:
#                        log.load()._apply(self)
#                    except EOFError, ex:
#                        break

    def _loadState(self):
        """
            Loads all checkpointed state.  Returns true if it was loaded, 
            false if none was found.
        """
        # load the big state file, if it exists
        stateFileName = os.path.join(self._path, 'State')
        try:
            fp = self._filer.openStateFile(stateFileName, Filer.READONLY)
            self._tables = Unpickler(fp).load()

            # restore the reference to the environment
            for table in self._tables.values():
                table._env = self

            return True

        except IOError:
            # create an empty database
            self._tables = {}
            return False

    def __init__(self, path, 
                 maxTxnFileSize = DEFAULT_MAX_TXN_FILE_SIZE,
                 filer = None,
                 checkpointCallback = None
                 ):

        if filer is None:
            filer = DefaultFiler(path, path)
        self._path = path
        self.__maxTxnFileSize = maxTxnFileSize
        self._filer = filer
        self.__checkpointCallback = checkpointCallback

        # database-wide lock
        self._lk = SafeLock()
        
        # checkpointing lock
        self.__checkpointLock = threading.Lock()

        # transaction used for all cursor operations
        self.__tmpTxn = None

        # make sure the directory exists
        self._filer.init()

        # make sure that the lock files exists (XXX race condition, also 
        # consider adding this to Filer)
        locks = []
        for lockFileName in  ('lock', 'checkpointLock'):
            lockFile = os.path.join(self._path, 'lock')
            if not os.path.exists(lockFile):
                locks.append(file(lockFile, 'w'))
            else:
                locks.append(file(lockFile, 'r+'))
        self.__lock, self.__checkpointLockFile = locks

        # get the lock, but don't try to read from the logfile yet
        self._acquireLock(update = False)

        try:
            # see if there is a lastlog file
            try:
                lastlogFp = \
                    self._filer.openStateFile('lastlog', Filer.READONLY)
                lastlog = lastlogFp.read()
                
                # if the lastlog file starts with a plus sign, we were 
                # checkpointing and the database could be in an inconsistent 
                # state.  A minus sign preceeds the last logfile that was read 
                # prior to completely checkpoointing.
                if lastlog.startswith('+'):
                    raise Exception('A program using the database appears to '
                                     'have terminated during the last '
                                     'checkpointing operation: if you have all '
                                     'of your logfiles, restore them to the '
                                     'database directory (%s) and edit %s to '
                                     'contain -%s' % (path, lastlog, 
                                                      LOG_PATTERN % 1
                                                      )
                                    )
                    
                elif lastlog.startswith('-'):
                    # Use the last logfile.
                    lastlog = lastlog[1:]
            except IOError, ex:
                # there is no lastlog file.  If there are state files, we will
                # assume that state is gone and attempt a recovery from 
                # logfiles. Otherwise we assume that this is just an old 
                # instance of the database and make the last logfile current.
                lastlog = None

            # load the current state
            if not self._loadState():
                # NO STATE.  Use an empty string to force a load of all of 
                # the logfiles.
                lastlog = ''
            
            # if there is a lastlog, load all of the logfiles after the 
            # lastlog except for the final, active one.
            logfiles = self._filer.getAllLogfiles()
            if lastlog is not None:
                logfile = None
                loadedLogs = False
                for logfile in logfiles[:-1]:
                    if logfile <= lastlog:
                        continue
                    print 'loading log %s' % logfile
                    self.__log = self._filer.openLogfile(logfile, 
                                                         Filer.READONLY
                                                         )
                    self.__readLog()
                    loadedLogs = True

                # if there were logfiles, checkpoint.
                if loadedLogs:
                    self.__checkpointAndCallback()

            # load the last logfile, if it exists.
            if logfiles:
                self.__log = \
                    self._filer.openLogfile(logfiles[-1], Filer.WRITABLE)
                self.__readLog()
            else:
                self.__openNewLog()
        finally:
            self._releaseLock()

    def __getLastLog(self):
        # find the last logfile - MUST HAVE THE DIRECTORY LOCK
        lastLog = None
        for filename in self._filer.getAllLogfiles():
            lastLog = filename
        return lastLog

    def __openNewLog(self):
        # open a new logfile - MUST HAVE THE DIRECTORY LOCK
        lastLog = self.__getLastLog()
        if lastLog:
            counter = int(lastLog.split('.')[-1]) + 1
        else:
            counter = 1
        filename = LOG_PATTERN % counter
        self.__log = self._filer.openLogfile(filename, Filer.CREATE)

    def _createBTreeDB(self, name, allowDups):
        return MemKeyDB(self, name, btree.BTree())

    def __getBTreeDB(self, txn, name, allowDups):
        txn._add(_CreateBTreeDBAction(name, allowDups))
        self._tables[name] = db = self._createBTreeDB(name, allowDups)
        return db

    def getBTreeDB(self, name, allowDups = False, txn = None):

        # first check the cache
        if self._tables.has_key(name):
            db = self._tables[name]
            if not isinstance(db._impl, btree.BTree):
                raise TypeError('%s is not a BTreeDB' % name)
            return db

        # create it transactionally
        return _needTxn(self.__getBTreeDB, self, txn, name, allowDups)
    
    def _createBListDB(self, name):
        return MemListDB(self, name, btree.BList())
    
    def __getBListDB(self, txn, name):
        txn._add(_CreateBListDBAction(name))
        self._tables[name] = db = self._createBListDB(name)
        return db
    
    def getBListDB(self, name, txn = None):
        
        # first check the cache
        if self._tables.has_key(name):
            db = self._tables[name]
            assert isinstance(db._impl, btree.BList)
            return db

        # create it transactionally
        return _needTxn(self.__getBListDB, self, txn, name)

    def _createHashDB(self, name, allowDups):
        return MemHashDB(self, name)

    def getHashDB(self, name, allowDups = False):

        # check the cache
        if self._tables.has_key(name):
            db = self._tables[name]
            assert isinstance(db, MemHashDB)
            return db

        # create it transactionally
        txn = MemTxn(self)
        try:
            txn._add(_CreateHashDBAction(name, allowDups))
            self._tables[name] = db = self._createHashDB(name, allowDups)
            txn.commit()
            txn = None
        finally:
            if txn: txn.abort()

        return db

    def getTxn(self):
        return MemTxn(self)

    def __runCheckpointCallback(self):
        try:
            self.__checkpointCallback(self._filer)
        finally:
            self._releaseCheckpointLock()

    def __writeLastLog(self, contents):
        lastlog = self._filer.openStateFile('lastlog', Filer.CREATE)
        lastlog.write(contents)
        lastlog.close()

    def __checkpointAndCallback(self):
        """
            Acquire the checkpoint lock, checkpoint and initiate the callback 
            (if there is one) in another thread).
        """
        checkpointing = False
        try:
            # get the checkpointing lock and set a flag to indicate to cleanp 
            # whether we are currently checkpointing.
            self._acquireCheckpointLock()
            checkpointing = True
            
            # store the lastlog file with an indicator that we're checkpointing
            lastLogName = self.__getLastLog()
            self.__writeLastLog('+' + lastLogName)

            self._checkpointTables()
            
            # rewrite lastlog to indicate that we have a complete state.
            self.__writeLastLog('-' + lastLogName)
            
            # if the user has provided a checkpointing callback run it in 
            # another thread and let it take over the checkpointing lock so 
            # we don't checkpoint again while the application layer is doing 
            # its thing.
            if self.__checkpointCallback:
                checkpointing = False
                thrd = threading.Thread(name = 'Checkpoint Callback',
                                        target = self.__runCheckpointCallback
                                        )
                thrd.start()
        finally:
            if checkpointing:
                self._releaseCheckpointLock()

    def checkpoint(self):
        self._acquireLock()
        try:
            self.__checkpointAndCallback()
            self.__openNewLog()
        finally:
            self._releaseLock()
    
    def close(self):
        self._acquireLock()
        try:
            for table in self._tables.values():
                table.close()
            
            self.__log.close()
            self.__log = None
        finally:
            self._releaseLock()
    
    def _checkpointTables(self):
        stateFileName = os.path.join(self._path, 'State')
        stateFile = self._filer.openStateFile(stateFileName, Filer.READONLY)
        Pickler(stateFile).dump(self._tables)

    def startDeadlockThread(self, deadlockCheckInterval, noticeCallback):
        pass

    def _acquireLock(self, update = True):
        # use both a mutex and a file system lock - the first should perform
        # better for locking between threads (and should also give us a nice
        # sane error if we double-lock), the second is essential for locking
        # between processes
        self._lk.acquire()
        fcntl.lockf(self.__lock.fileno(), fcntl.LOCK_EX)

        # read any transactions that we might have missed
        if update: self.__readLog()

    def _releaseLock(self):
        fcntl.lockf(self.__lock.fileno(), fcntl.LOCK_UN)
        self._lk.release()
    
    def _acquireCheckpointLock(self):
        self.__checkpointLock.acquire()
        fcntl.lockf(self.__checkpointLockFile, fcntl.LOCK_EX)
    
    def _releaseCheckpointLock(self):
        fcntl.lockf(self.__checkpointLockFile, fcntl.LOCK_UN)
        self.__checkpointLock.release()

    def _write(self, txn):
        self.__dumpTxn(txn)
        #self.__log.dump(txn)
        self.__log.flush()

    def apply(self, txn):
        assert txn._deferred(), 'Applying a transaction that was not deferred.'
        self._acquireLock()
        try:
            self._write(txn)
            txn._apply(self)
        finally:
            self._releaseLock()

    def _getTmpTxn(self):
        txn = self.__tmpTxn and self.__tmpTxn()
        if not txn:
            txn = MemTxn(self)
            self.__tmpTxn = ref(txn)
        return txn
    
    def _removeTable(self, name):
        del self._tables[name]

class MemCursor(Cursor):
    """
        Memory-oriented implementation of a database cursor.
    """

    def __init__(self, db, state, txn):
        self.__db = db
        self._state = state
        self.__txn = txn
        txn._addCursor(self)

    def __del__(self):
        self.close()

    def close(self):
        if self.__db is not None:
            self.__txn._removeCursor(self)
            self.__db = self._state = self.__txn = None

    def setToFirst(self):
        assert self.__db is not None
        self._state.setToFirst()

    def setToLast(self):
        assert self.__db is not None
        self._state.setToLast()

    def setToKey(self, key, exact = False):
        # XXX fix "exact"
        assert self.__db is not None
        self._state.setToKey(key)

    def __iter__(self):
        return self

class MemSeqCursor(MemCursor):
    """
        Cursor for sequential types that only returns values.
    """
    def next(self):
        if not self._state.initialized: self._state.setToFirst()
        return self._state.next()
    
    def prev(self):
        return self._state.prev()
    
class MemKeyValCursor(MemCursor):
    """
        Cursor for key databases that returns both the key and the value.
    """

    def next(self):
        if not self._state.initialized: self._state.setToFirst()
        key = self._state.getKey()
        return key, self._state.next()

    def prev(self):
        val = self._state.prev()
        return self._state.getKey(), val

class MemTxnBase(Txn):

    def __init__(self, env):
        self._env = env
        self._actions = []
        self.__cursors = WeakKeyDictionary()
        self.__closeCallbacks = []
        self.annotations = {}

    def __del__(self):
        try:
            if self._env:
                self.abort()
        except AttributeError:
            print >>sys.stderr, \
                'Error in MemTxn destructor: this is most likely because ' \
                'of an error unpickling one of your objects.'
            pass

    def _addCursor(self, cursor):
        self.__cursors[cursor] = True

    def _removeCursor(self, cursor):
        if self.__cursors.has_key(cursor):
            del self.__cursors[cursor]

    def _add(self, action):
        self._actions.append(action)

    def __setstate__(self, d):
        self.__dict__ = d
        self._env = None
        self.__cursors = WeakKeyDictionary()
        self.__closeCallbacks = []
        if not 'annotations' in d:
            self.annotations = {}
        
        # support older transactions
        if '_MemTxn__actions' in d:
            self._actions = self._MemTxn__actions
            del self._MemTxn__actions

    def __getstate__(self):
        d = self.__dict__.copy()
        del d['_env']
        del d['_MemTxnBase__cursors']
        del d['_MemTxnBase__closeCallbacks']
        
        # remove the annotations dictionary if it is empty to save space.
        if not self.annotations:
            del d['annotations']
        return d

    def _releaseLock(self):
        """
            Derived classes may override to release the environment lock on 
            commit or abort (from commonClose).
        """
        pass

    def addCloseCallback(self, callback):
        """
            Register a function to be called when the transaction is committed 
            or aborted.

            parms:
                callback: [callable<MemTxn>]
        """
        self.__closeCallbacks.append(callback)

    def __commonClose(self):
        "Clean-up function to be called both by commit and abort"

        try:
            if self.__cursors:
                # close all of the cursors, we have to use "keys()" to collect
                # them instead of straight out iteration because Cursor.close()
                # removes the cursor from the __cursors dictionary and we would
                # invalidate the iterator.
                for cur in self.__cursors.keys():
                    cur.close()
        finally:
            # if there's a screw-up with closing the cursor, we still need to
            # release the lock and release the environment
            self._actions = None
            if self._env: self._releaseLock()

            # disconnect so we don't try to do this again
            self._env = None
            
            # notify the interested
            for callback in self.__closeCallbacks:
                try:
                    callback(self)
                except:
                    print 'error calling transaction close callback %s:' % \
                          callback
                    traceback.print_exc()
    
    def commit(self):
        try:
            self._commit()
        finally:
            self.__commonClose()

    def abort(self):
        # unwind all of the actions that we performed
        try:
            if self._actions:
                self._actions.reverse()
                for action in self._actions:
                    action.undoIt(self._env)
        finally:
            self.__commonClose()

    def defer(self):
        # do all of the abort stuff, but restore the action list (make a copy 
        # of it because abort() reverses it)
        actions = list(self._actions)
        self.abort()
        self._actions = actions
    
    def _deferred(self):
        "Returns true if the transaction has been deferred."
        return self._actions is not None and self._env is None
    
    def isActiveOn(self, env):
        """
            Returns true if the transaction is active in the specfied 
            environment.
        """
        return env is self._env

    def dump(self, out, indent = ''):
        out.write(indent + 'Txn {\n')
        if self.annotations:
            out.write(indent + '  Annotations {\n')
            for ann, val in self.annotations.iteritems():
                out.write('%s    %s: %s\n' % (indent, repr(ann), repr(val)))
            out.write(indent + '  }\n')
        for act in self._actions:
            act.dump(out, indent + '  ')
        out.write(indent + '}\n')

    def _apply(self, env):
        # apply all actions to the database
        # this used to live in MemTxn, but now SubTxn's can also be applied if 
        # they are deferred().
        applied = []
        try:
            for action in self._actions:
                action.doIt(env)
                applied.append(action)
        except:
            # if we failed to apply some of our actions, undo the other 
            # actions.
            applied.reverse()
            for action in applied:
                action.undoIt(env)
            raise
        self._actions = None

class MemTxn(MemTxnBase):
    """
        Transaction implementation for memory databases.
    """

    def __init__(self, env):
        self._env = None
        env._acquireLock()
        MemTxnBase.__init__(self, env)

    def _commit(self):
        # if there are any actions, commit the transaction to the logfile
        if self._actions:
            self._env._write(self)

    def _releaseLock(self):
        self._env._releaseLock()

class SubTxn(MemTxnBase):
    """
        A sub-transaction.
        
        Public-vars:
            parent: [MemTxnBase] the parent transaction.
    """
    
    def __init__(self, env, parent):
        """
            parms:
                env: [Env]
                parent: [MemTxn or SubTxn] the parent transaction.
        """
        MemTxnBase.__init__(self, env)
        self.parent = parent
    
    def _commit(self):
        for action in self._actions:
            self.parent._add(action)    

class KeyVal(object):

    # define two "virtual" methods and use lambda to associate them with a 
    # property (otherwise they don't get virtualized)
    def _getVal(self): return self._val
    def _setVal(self, val): self._val = val
    val = property(lambda self: self._getVal(), 
                   lambda self, val: self._setVal(val)
                   )

    def __init__(self, key, val):
        self.key = key
        self._val = val
        self._maxKey = key

    def dump(self, indenter):
        print indenter, 'key = %s, val = %s' % \
            (repr(self.key), repr(self.val))
    
    def _markDeleted(self):
        "Allows derived classes to mark the node for deletion"
        pass

class MemKeyDB(KeyDB):
    """
        Abstract implementation of KeyDB which defines all of the standard
        functions in terms of the hidden implementation functions.
    """

    def __init__(self, env, name, impl):
        self._env = env
        self._name = name
        self._impl = impl

    def __getstate__(self):
        dict = self.__dict__.copy()
        del dict['_env']
        return dict

    def __set(self, txn, key, val):
        cur = self._impl.cursor().setToKey(key)
        
        # store the current value
        gotVal = cur.found()
        if gotVal:
            oldVal = cur.getVal()
        else:
            oldVal = None
            
        # set the new value
        cur.setValue(val)

        # store the action
        act = _ReplaceAction(self._name, key, val, oldVal, gotVal)
        txn._add(act)
    
    def set(self, txn, key, val):
        return _needTxn(self.__set, self._env, txn, key, val)
    
    def _replace(self, key, val):
        cur = self._impl.cursor().setToKey(key)
        cur.setValue(val)
        
    def _insert(self, key, val):
        cur = self._impl.cursor().setToKey(key)
        cur.insert(val)

    def __getitem__(self, key):
        txn = self._env._getTmpTxn()
        cur = self._impl.cursor().setToKey(key)
        del txn
        if not cur.found():
            raise KeyError(key)
        else:
            return cur.getVal()

    def get(self, key, default = None, txn = None):
        if not txn:
            txn = myTxn = self._env._getTmpTxn()
        cur = self._impl.cursor().setToKey(key)
        if not cur.found():
            return default
        else:
            return cur.getVal()

    def has_key(self, key):
        txn = self._env._getTmpTxn()
        cur = self._impl.cursor().setToKey(key)
        return cur.found()

    def __delete(self, txn, key):
        # delete the key from the implementation
        val = self._delete(key)
        
        # add it to the transaction
        txn._add(_DeleteAction(self._name, key, val))

    def delete(self, txn, key):
        _needTxn(self.__delete, self._env, txn, key)
    
    def _delete(self, key):
        cur = self._impl.cursor().setToKey(key)
        if not cur.found(): raise KeyError(key)
        val = cur.getVal()
        cur.delete()
        return val

    def __delitem__(self, key):
        txn = MemTxn(self._env)
        self.delete(txn, key)
        txn.commit()

    def __setitem__(self, key, val):
        txn = MemTxn(self._env)
        try:
            self.set(txn, key, val)
            txn.commit()
            txn = None
        finally:
            if txn: txn.abort()

    def getTxn(self):
        return self._env.getTxn()

    def cursor(self, txn = None):
        # make sure that we do this in the scope of a transaction
        if not txn:
            txn = self._env._getTmpTxn()

        return self._createCursor(self._impl.cursor(), txn)
    
    def _createCursor(self, implCursor, txn):
        return MemKeyValCursor(self, implCursor, txn)
    
    def close(self):
        self._env._removeTable(self._name)
        self._env = None
    
    def checkpoint(self):
        self._impl.checkpoint()
    
    def dump(self):
        self._impl.dump()

class MemListDB(MemKeyDB):  # XXX should be derived from ListDB?
    
    def __append(self, txn, val):
        txn._add(_AppendAction(self._name, val))
        self._append(val)

    def append(self, val, txn = None):
        _needTxn(self.__append, self._env, txn, val)
    
    def __insert(self, txn, index, val):
        txn._add(_InsertAction(self._name, index, val, None, False))
        self._insert(index, val)
    
    def insert(self, index, val, txn = None):
        _needTxn(self.__insert, self._env, txn, index, val)

    def __pop(self, txn, index):
        result = self._impl.pop(index)
        txn._add(_PopAction(self._name, index, result))
        return result
    
    def pop(self, index = -1, txn = None):
        """
            Pop an object off the list db.
            
            parms:
                index: [int] The object index.
                txn: [@MemTxn or None] the transaction
        """
        return _needTxn(self.__pop, self._env, txn, index)

    def __len(self, txn):
        return len(self._impl)

    def len(self, txn):
        """
            Returns the length of the table in a transactional context.
            
            parms:
                txn: [@MemTxn]
        """
        return _needTxn(self.__len, self._env, txn)

    def __len__(self):
        """
            Implements the len() interface.  Returns the number of elements in 
            the table.
        """
        # It's painful that we have to create a transaction for this, but I'm 
        # concerned about the possibility that this operation might not always 
        # be atomic.
        return _needTxn(self.__len, self._env, None)
    
    def _createCursor(self, implCursor, txn):
        return MemSeqCursor(self, implCursor, txn)
    
    def _pop(self, index = -1):
        return self._impl.pop(index)
    
    def _append(self, val):
        self._impl.append(val)
    
class MemHashDB(MemKeyDB):
    """
        Hash table implementation.
    """

    def __init__(self, env, name):
        MemKeyDB.__init__(self, env, name, None)
        self.__dict = {}

    def _insert(self, key, val):
        self.__dict[key] = val

    def _delete(self, key):
        del self.__dict[key]

    def _lookUp(self, key):
        if self.__dict.has_key(key):
            return self.__dict[key], True
        else:
            return None, False

    # this is badly broken - but who cares about this for a hash db anyhow?
    def _find(self, key, exact): return self.__dict.iteritems()
    def _first(self): return self.__dict.iteritems()
    def _last(self): return self.__dict.iteritems()

    def _next(self, state):
        return (state,) + state.next()

    def cursor(self, txn = None):
        if not txn:
            txn = self._env._getTmpTxn()
        return MemCursor(self, self.__dict.iteritems(), txn)

if __name__ == '__main__':

    # this stuff needs to be done in a function because if we let it get
    # cleaned up in the global destructors we get an exception that appears to
    # be related to the order of destruction of the system objects
    def testfunc():
        print 'running test suite'
        import shutil
        if os.path.exists('utdata'): shutil.rmtree('utdata')
        env = MemEnv('utdata')
        db = env.getBTreeDB('First')

        # check the empty cursor condition
        print 'checking empty cursors'
        count = 0
        for key, val in db.cursor():
            count += 1
        assert count == 0

        # put some items in the database, make sure we get them all out
        print 'testing insertions'
        elems = [
            ('hello', 'world'),
            ('first', 100),
            ('second', 200),
            ('force tree', 300),
        ]
        elems.sort()
        sortedElems = elems[:]
        for elem in elems:
            db[elem[0]] = elem[1]

        # make sure that everything comes back out in order
        print 'checking ordered traversal'
        cur = db.cursor()
        for key, val in cur:
            assert (key, val) == elems.pop(0)
        assert elems == []

        # check reverse iteration
        print 'checking reverse iteration'
        elems = sortedElems[:]
        elems.reverse()
        cur = db.cursor()
        cur.setToLast()
        results = []
        for key, val in ReverseCursor(cur):
            assert (key, val) == elems.pop(0)

        # check cursor reversibility
        print 'checking reversibility'
        cur = db.cursor()
        key, val == cur.next()
        assert (key, val) == cur.prev()
        assert (key, val) == cur.next()
        cur.close()

        print 'checking for overwrites'
        db['hello'] = 'mother'
        db['hello'] = 'world'
        cur = db.cursor()
        for key, val in sortedElems:
            assert (key, val) == cur.next()
        try:
            cur.next()
            assert 0, "more stuff than expected in the db"
        except StopIteration:
            pass
        cur.close()

        # check the "getitem" functionality
        assert env.getBTreeDB('First')['hello'] == 'world'

        print 'checking hash databases'
        hdb = env.getHashDB('Second')
        helems = [('check', 100), ('out', 300), ('this', 200)]
        for key, val in helems:
            hdb[key] = val
        results = []
        for key, val in hdb.cursor():
            results.append((key, val))
        results.sort()
        assert results == [('check', 100), ('out', 300), ('this', 200)]

        print 'checking btrees with duplicate keys'
        ddbElems = [
            ('first', 100), 
            ('second', 200), 
            ('second', 300), 
            ('third', 400)
        ]
        ddb = env.getBTreeDB('BTreeWithDups', allowDups = True)
        for key, val in ddbElems:
            ddb[key] = val
        ddbElems.sort()
        cur = ddb.cursor()
        for key, val in ddbElems:
            #k,v = cur.next()
            assert (key, val) == cur.next()
        try:
            cur.next()
            assert 0, 'too many elements'
        except StopIteration:
            pass

        # now verify persistence
        print 'verifying persistence'
        del db, env, hdb
        env = MemEnv('utdata')
        db = env.getBTreeDB('First')
        assert db['hello'] == 'world'
        elems = sortedElems[:]
        for key, val in db.cursor():
            assert (key, val) == elems.pop(0)
        assert elems == []

        print 'verifying hash persistence'
        results = []
        hdb = env.getHashDB('Second')
        for key, val in hdb.cursor():
            results.append((key, val))
        results.sort()
        assert results == [('check', 100), ('out', 300), ('this', 200)]

        print 'hash deletion'
        txn = hdb.getTxn()
        hdb.set(txn, 'delete', 'me')
        hdb.delete(txn, 'delete')
        txn.commit()
        results = []
        for key, val in hdb.cursor():
            results.append((key, val))
        results.sort()
        assert results == [('check', 100), ('out', 300), ('this', 200)]

        # check for removal
        print 'verifying removal'
        del db[sortedElems[0][0]]
        try:
            db[sortedElems[0][0]]
            assert 0, 'removal failed'
        except KeyError, ex:
            assert ex[0] == sortedElems[0][0]
        # restore it
        db[sortedElems[0][0]] = sortedElems[0][1]

        # test checkpoints
        print 'testing checkpoints'
        env.checkpoint()

        del env, db, hdb
        env = MemEnv('utdata')
        db = env.getBTreeDB('First')
        results = []
        for key, val in db.cursor():
            results.append((key, val))
        assert results == sortedElems

        hdb = env.getHashDB('Second')
        results = []
        for key, val in hdb.cursor():
            results.append((key, val))
        results.sort()
        assert results == helems

        print 'double locking'
        try:
            txn = env.getTxn()
            txn2 = env.getTxn()
            assert 0, 'double locks are not giving errors'
        except DoubleLockError:
            pass


        print 'all tests passed'
    testfunc()

# garrghhh... cyclic links
import btree

# vim:set sw=4: